// Copyright DustedPixels.com 2008. All rights reserved.

package com.dustedpixels.asm.utils;

import org.objectweb.asm.Label;
import org.objectweb.asm.MethodAdapter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

import com.dustedpixels.asm.InsnAdviceAdapter;


/**
 * @author micapolos@gmail.com (Michal Pociecha-Los)
 */
public final class ReturnRemover {
  /**
   * Removes all RETURN instructions from visited method modifying method
   * bytecode in such way, that the only place the control is left from the
   * method is its end, and writes the result into underlying output.
   * 
   * <p>
   * This adapter is one of the key adapter preparing method bytecode to be inlined.
   * 
   * @param out return remover output
   * @return return return removing visitor
   */
  public static MethodVisitor newOn(MethodVisitor out) {
    return new ReturnRemoverImpl(out);
  }
  
  /**
   * Implementation of return removing visitor.
   */
  private final static class ReturnRemoverImpl extends InsnAdviceAdapter {

    private Label gotoLabel = null;
    private boolean returnDetected = false;
    
    public ReturnRemoverImpl(MethodVisitor out) {
      super(out);
    }

    @Override
    public void visitInsn(int opcode) {
      switch (opcode) {
        case Opcodes.IRETURN:
        case Opcodes.LRETURN:
        case Opcodes.FRETURN:
        case Opcodes.DRETURN:
        case Opcodes.RETURN:
          returnDetected = true;
          break;
        default:
          super.visitInsn(opcode);
          break;
      }
    }
    
    @Override
    public void visitMaxs(int maxStack, int maxLocals) {
      if (gotoLabel != null) {
        mv.visitLabel(gotoLabel);
      }
      mv.visitMaxs(maxStack, maxLocals);
    }

    @Override
    public void visitLabel(Label label) {
      if (returnDetected) {
        if (gotoLabel == null) {
          gotoLabel = new Label();
        }
        mv.visitJumpInsn(Opcodes.GOTO, gotoLabel);
        returnDetected = false;
      }
      mv.visitLabel(label);
    }

    @Override
    protected void beforeInsn() {
      if (returnDetected) {
        if (gotoLabel == null) {
          gotoLabel = new Label();
        }
        mv.visitJumpInsn(Opcodes.GOTO, gotoLabel);
        returnDetected = false;
      }
    }
  }
}
